#include "nBodySimulator.h"

// Initialize the simulation by reading in the parameters and the
// initial setup.
nBodySimulator::nBodySimulator(string parfn, string initfn)
{
   params  = new ParamHolder(parfn);
   holder = new OctreeHolder(initfn, params);
}

// Run the simulation using the specified parameters.
void nBodySimulator::run()
{
   stepCounter =0;

   double dt;
   OutputNumber = 0;

   t = params -> StartTime;

   double nextOutput = 0.0;

   // Advance by timesteps until the stop time is reached.
   while (t < params -> StopTime) {

      // Try the default time step.
      dt = params -> dt;

      if (nextOutput > params -> StopTime)
         nextOutput = params -> StopTime;

      // If an output occurs before that, run up to that output.
      // and write out the state of the simulation.
      if (nextOutput < t+dt) {
         dt = nextOutput - t;

	 advanceSim(dt);
	 t = nextOutput;
	 nextOutput += params -> OutputDt;

         // Figure out the name of the output file.
	 stringstream ss;

	 ss << "data/" << params -> OutputPrefix <<  setfill('0') << setw(params -> ZerosInOutputName) << OutputNumber;
	 string ofName = ss.str();

	 holder -> writeToFile(ofName);

	 OutputNumber++;
	 cout << t << endl;
         }

      // Otherwise, just run one timestep.
      else {
         advanceSim(dt);
	 t += dt;
         }

      }

      cout << "Steps taken: " << stepCounter << endl;
}

// This is what actually update the positions and velocities.
// Uses a 2nd order KDK leapfrog algorithm.
void nBodySimulator::advanceSim(double dt)
{
   int n = holder -> size();

   vec* accels = new vec[n];
   int* adtFrac = new int[n];

   double magAccel, adt;
   int fracDTs = 1; // How many fractional timesteps?

   intCounter counter;

   // 1. Compute the acceleration at the beginning of the timestep for
   // each body.
   for (int i = 0; i < n; i++)
      accels[i] = holder -> accel(i);

   // 2. Figure out the fraction of the maximum timestep to use for
   // each body.
   for (int i = 0; i < n; i++) {
      magAccel = sqrt(pow(accels[i][X], 2.0) + pow(accels[i][X], 2.0) + pow(accels[i][X], 2.0) );

      adt = sqrt(2.0 * params -> TimeStepTol * params -> SofteningParameter / magAccel);

      if (adt > dt) // Never use more than the max timestep!
         adt = dt;

      adtFrac[i] = (int)(pow(2.0, ceil(-log(adt/dt)/log(2.0))));

      if (adtFrac[i] < 1)
         adtFrac[i] = 1;

      if (adtFrac[i] > fracDTs)
         fracDTs = adtFrac[i];

      counter.increment(adtFrac[i]);
      }

   //cout << "Start: " << endl;
   //counter.printCounter();
   //cout << endl;


   // 3. Using the smallest fraction of a timestep, update each
   // particle.

   for (int f = 0; f < fracDTs; f++) {
      double fdt = dt / (double)fracDTs;

      // 3a. Kick each particle that needs to be kicked to t+0.5*fdt
      for (int i = 0; i < n; i++) { 
         Body* b = holder -> getBody(i);

         if (f % (int)(fracDTs/adtFrac[i]) == 0) {
            b -> vel = b -> vel + accels[i] * (fdt * (double)fracDTs/(double)adtFrac[i]) * 0.5;
	    }

         }
     
      // 3b. Drift all particles to t+fdt
      for (int i = 0; i < n; i++) {
         Body* b = holder -> getBody(i);
         b -> pos = b->pos + b->vel * fdt;
	 }
 
      stepCounter++;
      holder -> rebuildTree();

      // 3c. Kick the particles that need to be kicked at t+fdt.
      for (int i = 0; i < n; i++)
         accels[i] = holder -> accel(i);

      for (int i = 0; i < n; i++) {
         Body* b = holder -> getBody(i);

         if ((f+1) % (fracDTs/adtFrac[i]) == 0)
            b -> vel = b -> vel + accels[i] * (fdt * (double)(fracDTs/adtFrac[i]))* 0.5;
         }

      // 3d. Check if any particles need to have a smaller timestep.
      for (int i = 0; i < n; i++) {
         magAccel = sqrt(pow(accels[i][X], 2.0) + pow(accels[i][X], 2.0) + pow(accels[i][X], 2.0) );
         adt = sqrt(2.0 * params -> TimeStepTol * params -> SofteningParameter / magAccel);

         int den = (int)(pow(2.0, ceil(-log(adt/dt)/log(2.0))));

	 if (den > adtFrac[i]) {
	    if (den > fracDTs) {
	       for(int j = 0; j < n; j++) {
	          if (j != i)
	             adtFrac[j] *= (den/adtFrac[i]);
	          }

	       fracDTs *= (den/adtFrac[i]);
	       f *= (int)(den/adtFrac[i]);
	       }
	    counter.decrement(adtFrac[i]);
	    adtFrac[i] = den;
	    counter.increment(adtFrac[i]);
	    }

         }

      }

   //cout << "End: " << endl;
  // counter.printCounter();
  // cout << endl;

   delete [] accels;
}
